#include "BTreeNode.h"

using namespace std;


BTLeafNode::BTLeafNode()
{
    numKeys = (int*)buffer;
    nextPageId = (int*)buffer + 1;
    leafPairs = (BTLeafPair*) ((int*)buffer + 2);
}
/*
 * Read the content of the node from the page pid in the PageFile pf.
 * @param pid[IN] the PageId to read
 * @param pf[IN] PageFile to read from
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::read(PageId pid, const PageFile& pf)
{
    RC rc = 0;
    if (rc = pf.read(pid, buffer) < 0) {
        fprintf(stderr, "Error: Could not read from PageFile\n");
        return rc;
    }

    return rc; 
}
    
/*
 * Write the content of the node to the page pid in the PageFile pf.
 * @param pid[IN] the PageId to write to
 * @param pf[IN] PageFile to write to
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::write(PageId pid, PageFile& pf)
{
    RC rc = 0;

    // Write buffer to PageFile
    if (rc = pf.write(pid, buffer)) {
        fprintf(stderr, "Error: Could not write to PageFile\n");
    }
    return rc;
}

/*
 * Return the number of keys stored in the node.
 * @return the number of keys in the node
 */
int BTLeafNode::getKeyCount()
{ 
    return *numKeys; 
}

/*
 * Insert a (key, rid) pair to the node.
 * @param key[IN] the key to insert
 * @param rid[IN] the RecordId to insert
 * @return 0 if successful. Return an error code if the node is full.
 */
RC BTLeafNode::insert(int key, const RecordId& rid)
{ 
    if (*numKeys >= MAX_LEAF_KEYS) {
        fprintf(stderr, "Error: Could not insert key to page, node full\n");
        return RC_NODE_FULL;
    }

    // Find where to insert this new pair
    int i = 0;
    while (key > leafPairs[i].key && i < *numKeys) {
        i++;
    }
    
    // Shift remaining pairs to make room for new pair
    for (int j = *numKeys; j >= i; --j) {
        leafPairs[j].key = leafPairs[j-1].key;
        leafPairs[j].rid.pid = leafPairs[j-1].rid.pid;
        leafPairs[j].rid.sid = leafPairs[j-1].rid.sid;
    }

    // Insert new pair
    leafPairs[i].key = key;
    leafPairs[i].rid.pid = rid.pid;
    leafPairs[i].rid.sid = rid.sid;

    // Increment numKeys
    *numKeys = *numKeys + 1;

    return 0;
}

/*
 * Insert the (key, rid) pair to the node
 * and split the node half and half with sibling.
 * The first key of the sibling node is returned in siblingKey.
 * @param key[IN] the key to insert.
 * @param rid[IN] the RecordId to insert.
 * @param sibling[IN] the sibling node to split with. This node MUST be EMPTY when this function is called.
 * @param siblingKey[OUT] the first key in the sibling node after split.
 * @return 0 if successful. Return an error code if there is an error.
 */
// Note: does not write either node to PageFile nor update nextPageId's
RC BTLeafNode::insertAndSplit(int key, const RecordId& rid, 
                              BTLeafNode& sibling, int& siblingKey)
{ 
    int newNumKeys = (*numKeys + 1) >> 1;
    int siblingNumKeys = ((*numKeys + 1) >> 1) + ((*numKeys + 1) % 2);

    // Set return value for sibling key
    siblingKey = leafPairs[0].key;

    // Add pairs to sibling
    int newCount = 0, oldCount = 0;
    while (newCount < siblingNumKeys) {
        if (key < leafPairs[oldCount].key) 
        {
            sibling.insert(key, rid);
            newCount++;
        } 
        else 
        {
            sibling.insert(leafPairs[oldCount].key, leafPairs[oldCount].rid);
            newCount++;
            oldCount++;
        }
    }

    // Shift remaining pairs to beginning of this leaf
    newCount = 0;
    while (newCount < newNumKeys) {
        if (key < leafPairs[oldCount].key) 
        {
            leafPairs[newCount].key = key;
            leafPairs[newCount].rid.pid = rid.pid;
            leafPairs[newCount].rid.sid = rid.sid;
            newCount++;
        } 
        else 
        {
            leafPairs[newCount].key = leafPairs[oldCount].key;
            leafPairs[newCount].rid.pid = leafPairs[oldCount].rid.pid;
            leafPairs[newCount].rid.sid = leafPairs[oldCount].rid.sid;
            newCount++;
            oldCount++;
        }
    }

    // Update numKeys
    *numKeys = newNumKeys;

    return 0; 
}

/*
 * Find the entry whose key value is larger than or equal to searchKey
 * and output the eid (entry number) whose key value >= searchKey.
 * Remeber that all keys inside a B+tree node should be kept sorted.
 * @param searchKey[IN] the key to search for
 * @param eid[OUT] the entry number that contains a key larger than or equalty to searchKey
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::locate(int searchKey, int& eid)
{ 
    eid = 0;
    while (leafPairs[eid].key < searchKey) 
    {
        eid++;
        if (eid == *numKeys) {
            fprintf(stderr, "Error: Could not locate entry in node\n");
            return RC_NO_SUCH_RECORD;
        }
    }
    return 0; 
}

/*
 * Read the (key, rid) pair from the eid entry.
 * @param eid[IN] the entry number to read the (key, rid) pair from
 * @param key[OUT] the key from the entry
 * @param rid[OUT] the RecordId from the entry
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::readEntry(int eid, int& key, RecordId& rid)
{ 
    if (eid >= *numKeys || eid < 0) {
        fprintf(stderr, "Error: Could not locate entry in node\n");
        return RC_NO_SUCH_RECORD;
    }
    
    key = leafPairs[eid].key;
    rid.pid = leafPairs[eid].rid.pid;
    rid.pid = leafPairs[eid].rid.sid;

    return 0; 
}

/*
 * Return the pid of the next slibling node.
 * @return the PageId of the next sibling node 
 */
PageId BTLeafNode::getNextNodePtr()
{ 
    return *nextPageId; 
}

/*
 * Set the pid of the next slibling node.
 * @param pid[IN] the PageId of the next sibling node 
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::setNextNodePtr(PageId pid)
{ 
    *nextPageId = pid;
    return 0; 
}

BTNonLeafNode::BTNonLeafNode()
{
    numKeys = (int*)buffer;
    keys = (int*)buffer + 1;
    pageIds = (int*)buffer + 1 + MAX_NONLEAF_KEYS;
}

/*
 * Read the content of the node from the page pid in the PageFile pf.
 * @param pid[IN] the PageId to read
 * @param pf[IN] PageFile to read from
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::read(PageId pid, const PageFile& pf)
{ 
    RC rc = 0;
    if (rc = pf.read(pid, buffer) < 0) {
        fprintf(stderr, "Error: Could not read from PageFile\n");
        return rc;
    }

    return 0; 
}

/*
 * Write the content of the node to the page pid in the PageFile pf.
 * @param pid[IN] the PageId to write to
 * @param pf[IN] PageFile to write to
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::write(PageId pid, PageFile& pf)
{ 
    RC rc = 0;

    if (rc = pf.write(pid, buffer) < 0) {
        fprintf(stderr, "Error: Could not read from PageFile\n");
        return rc;
    }

    return 0;  
}

/*
 * Return the number of keys stored in the node.
 * @return the number of keys in the node
 */
int BTNonLeafNode::getKeyCount()
{ 
    return *numKeys; 
}


/*
 * Insert a (key, pid) pair to the node.
 * @param key[IN] the key to insert
 * @param pid[IN] the PageId to insert
 * @return 0 if successful. Return an error code if the node is full.
 */
RC BTNonLeafNode::insert(int key, PageId pid)
{ 
    if (*numKeys >= MAX_NONLEAF_KEYS) 
    {
        fprintf(stderr, "Error: Could not insert key to page, node full\n");
        return RC_NODE_FULL;
    }

    // Find where to insert key and PageId
    int i = 0;
    while (i < *numKeys && key < keys[i])
    {
        i++;
    }

    // Shift remaining keys and PageId's down
    int newNumKeys = *numKeys + 1;
    for (int j = *numKeys; j >= i; --j)
    {
        keys[j+1] = keys[j];
        pageIds[j+2] = pageIds[j+1];
    }

    // Insert new key and PageId
    keys[i] = key;
    pageIds[i+1] = pid;

    // Increment numKeys
    *numKeys = *numKeys + 1;

    return 0; 
}

/*
 * Insert the (key, pid) pair to the node
 * and split the node half and half with sibling.
 * The middle key after the split is returned in midKey.
 * @param key[IN] the key to insert
 * @param pid[IN] the PageId to insert
 * @param sibling[IN] the sibling node to split with. This node MUST be empty when this function is called.
 * @param midKey[OUT] the key in the middle after the split. This key should be inserted to the parent node.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::insertAndSplit(int key, PageId pid, BTNonLeafNode& sibling, int& midKey)
{ 
    int newNumKeys = (*numKeys >> 1) + (*numKeys % 2);
    int siblingNumKeys = *numKeys  >> 1;

    // Find where to insert key and PageId
    int newKeyIndex = 0;
    while (newKeyIndex < newNumKeys && key < keys[newKeyIndex])
    {
        newKeyIndex++;
    }

    // Check if new key is midKey
    if (newKeyIndex == newNumKeys && key < keys[newKeyIndex])
    {
        midKey = key;

        // Initialize sibling with new pid and first of second half of keys
        sibling.initializeRoot(pid, keys[newNumKeys], pageIds[newNumKeys+1]);

        // Copy second half to sibling
        for (int j = newNumKeys+1; j < siblingNumKeys; ++j)
        {
            sibling.insert(keys[j], pageIds[j+1]);
        }
    }

    // Check if new key needs to be inserted in this node
    else if (newKeyIndex < newNumKeys)
    {
        midKey = keys[newNumKeys];

        // Copy second half of keys to sibling
        sibling.initializeRoot(pageIds[newNumKeys], keys[newNumKeys], pageIds[newNumKeys + 1]);
        for (int j = 0; j < siblingNumKeys; ++j)
        {
            sibling.insert(keys[newNumKeys + j], pageIds[newNumKeys + 1 + j]);
        }
    }

    // Otherwise new key goes in sibling
    else
    {
        midKey = keys[newNumKeys];

        // If new key is less than first of new sibling, initialize root with it
        if (key < keys[newNumKeys+1])
        {
            sibling.initializeRoot(pageIds[newNumKeys+1], key, pid);
            for (int j = 1; j < siblingNumKeys; ++j)
            {
                sibling.insert(keys[newNumKeys+j], pageIds[newNumKeys+1+j]);
            }
        }
        else
        {
            // Initialize root with first non-midKey
            sibling.initializeRoot(pageIds[newNumKeys+1], keys[newNumKeys+1], pageIds[newNumKeys+2]);

            // Copy over keys and PageIds, checking for where to insert new key
            int i = newNumKeys + 2;
            for (int siblingCount = 1; siblingCount < siblingNumKeys; ++siblingCount)
            {
                if (keys[i] > key)
                {
                    sibling.insert(key, pid);
                }
                else
                {
                    sibling.insert(keys[i], pageIds[i+1]);
                    i++;
                }
            }
        }
    }

    // Set new numKeys
    *numKeys = newNumKeys;

    return 0; 
}

/*
 * Given the searchKey, find the child-node pointer to follow and
 * output it in pid.
 * @param searchKey[IN] the searchKey that is being looked up.
 * @param pid[OUT] the pointer to the child node to follow.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::locateChildPtr(int searchKey, PageId& pid)
{ 
    int i = 0;
    while ((i < *numKeys + 1) && searchKey > keys[i])
    {
        i++;
    }
    pid = pageIds[i];
    return 0; 
}

/*
 * Initialize the root node with (pid1, key, pid2).
 * @param pid1[IN] the first PageId to insert
 * @param key[IN] the key that should be inserted between the two PageIds
 * @param pid2[IN] the PageId to insert behind the key
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::initializeRoot(PageId pid1, int key, PageId pid2)
{ 
    keys[0] = key;
    pageIds[0] = pid1;
    pageIds[1] = pid2;
    *numKeys = 1;
    return 0; 
}
